import os

os.environ["MUJOCO_GL"] = "egl"

import mujoco
import numpy as np
from mujoco import minimize
from scipy.spatial.transform import Rotation


class MujocoIK:
    """
    The radius is the weight of the orientation residual (in Lie algebra)
    The smaller the radius, the more the orientation residual is ignored.

    Forward functions in mujoco:

    mujoco.mj_kinematics(self.model, self.data):
        compute the kinematics of the model and update the data inplace

    mujoco.mj_forward(self.model, self.data)
        = {
        1. mj_kinematics(model, data),  # update the position and the rotations
        2. compute the constraint forces and the accelerations
    }

    mujoco.mj_step(self.model, self.data):
        = {
        1. mj_forward(model, data),  # update the position and the rotations, forces and accelerations
        2. Do the integration of the dynamics
    }

    In this case, only position and rotation are updated, so mj_kinematics is enough.

    Mujoco use quaternion as [w, x, y, z]
    pinocchio use quaternion as [x, y, z, w]

    The joint dimension is 8: 6 joints + 2 gripper angle
    """

    def __init__(self, radius=0.04, reg=0.01, max_iter=1000):
        self.model = mujoco.MjModel.from_xml_path("piper_description/mujoco_model/piper_description.xml")
        self.data = mujoco.MjData(self.model)

        # Reset the state to the "zero_position" keyframe.
        key = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_KEY, "zero_position")
        mujoco.mj_resetDataKeyframe(self.model, self.data, key)
        mujoco.mj_forward(self.model, self.data)

        # Bounds at the joint limits.
        self.bounds = [self.model.jnt_range[:, 0], self.model.jnt_range[:, 1]]

        # Inital guess is the 'home' keyframe if nothing is provided
        self.j_home = self.model.key("home").qpos
        self.radius = radius
        self.reg = reg  # apply different regularization to different joints
        self.max_iter = max_iter

    def ik_residual(self, joint, pos=None, quat=None, reg_target=None):
        """Residual for inverse kinematics.

        Args:
        joint: joint angles.
        pos: target position for the end effector.
        quat: target orientation for the end effector.
        radius: scaling of the 3D cross.

        Returns:
        The residual of the Inverse Kinematics task.
        """
        # Move the mocap body to the target
        tar_id = self.model.body("ik_target").mocapid
        self.data.mocap_pos[tar_id] = self.model.body("ik_target").pos if pos is None else pos
        self.data.mocap_quat[tar_id] = self.model.body("ik_target").quat if quat is None else quat

        # Set qpos, compute forward kinematics.
        res = []
        for i in range(joint.shape[1]):
            self.data.qpos = joint[:, i]
            mujoco.mj_kinematics(self.model, self.data)

            # Position residual: can also use Lie algebra to compute the residual
            res_pos = self.data.site("effector").xpos - self.data.site("ik_target").xpos
            # res_pos = self.data.site("effector").xpos - pos

            # Effector quat, use mju_mat2quat.
            effector_quat = np.empty(4)
            mujoco.mju_mat2Quat(effector_quat, self.data.site("effector").xmat)

            # Target quat, exploit the fact that the site is aligned with the body.
            target_quat = self.data.body("ik_target").xquat
            # target_quat = quat

            # Orientation residual: quaternion difference to lie algebra (w1, w2, w3), which defines the rotation axis
            # the rotation angle  theta^2 = w1^2 + w2^2 + w3^2
            res_lie = np.empty(3)  # Lie axis (w1, w2, w3)
            mujoco.mju_subQuat(res_lie, target_quat, effector_quat)
            res_lie *= self.radius

            # Regularization residual. To avoid the joint angles from being too far from the home position.
            # reg_target=x_prev is to regularize the joint angles to be close to the previous joint angles
            reg_target = self.model.key("home").qpos if reg_target is None else reg_target
            res_reg = self.reg * (joint[:, i] - reg_target)

            res_i = np.hstack((res_pos, res_lie, res_reg))
            res.append(np.atleast_2d(res_i).T)

        return np.hstack(res)

    def ik_jac(self, joint, res, pos=None, quat=None):
        """Analytic Jacobian of inverse kinematics residual

        Args:
        joint: joint angles.
        pos: target position for the end effector.
        quat: target orientation for the end effector.
        radius: scaling of the 3D cross.

        Returns:
        The Jacobian of the Inverse Kinematics task.
        """
        # least_squares() passes the value of the residual at x which is sometimes
        # useful, but we don't need it here.
        del res

        # Call mj_kinematics and mj_comPos (required for Jacobians).
        mujoco.mj_kinematics(self.model, self.data)
        mujoco.mj_comPos(self.model, self.data)

        # Get end-effector site Jacobian.
        jac_pos = np.empty((3, self.model.nv))
        jac_lie = np.empty((3, self.model.nv))
        mujoco.mj_jacSite(self.model, self.data, jac_pos, jac_lie, self.data.site("effector").id)

        effector_quat = np.empty(4)
        mujoco.mju_mat2Quat(effector_quat, self.data.site("effector").xmat)
        target_quat = self.data.body("ik_target").xquat
        # target_quat = quat
        # the 3x3 mjd_subQuat Jacobian (derivative of quaternion error)
        D_effector = np.empty((3, 3))  # ∂ω / ∂qb
        # Related to: R_error = R_target @ R_effector^-1, is the rotation that rotates the effector to the target
        # w = log(R_error), but here is the derivative of w
        # note that error != target_quat - effector_quat, because the error is in the tangent space of SO(3)
        # The perfect rotation is R_error = I, so w = 0
        mujoco.mjd_subQuat(target_quat, effector_quat, None, D_effector)

        # Transform to target frame and apply chain rule for error Jacobian
        # target_mat: = Rotation.from_quat(quat).as_matrix() = rotation matrix of the target
        target_rotation = self.data.site("ik_target").xmat.reshape(3, 3)
        # Rotate into target frame, multiply by subQuat Jacobian, scale by radius.
        # WHY: jac_lie: joints -> effector -> target frame.
        # D_effector = log(R_target @ R_effector^-1)： rotate from effector to target frame.
        # but in the mse loss, we use D_effector inverse, which is a rotation from target to effector frame.
        # Chain rule: ∂error/∂q = ∂error/∂target × ∂target /∂q
        #             ∂error/∂q = ∂error/∂target × ∂(target_rotation × effector) /∂q
        #             ∂error/∂q = ∂error/∂target × target_rotation × ∂(effector) /∂q
        jac_lie = self.radius * D_effector.T @ target_rotation.T @ jac_lie

        # Regularization Jacobian.
        jac_reg = self.reg * np.eye(self.model.nv)

        return np.vstack((jac_pos, jac_lie, jac_reg))

    def solve_ik_xquat(self, pos=None, quat=None, j0=None, x_prev=None):
        """
        Solve the IK problem.
        pos: target position for the end effector.
        quat: target orientation for the end effector.
        j0: initial guess of the joint angles. If None, use the home position.
        """
        if j0 is None:
            # initial guess is the home position
            j0 = self.j_home
        ik_target = lambda joint: self.ik_residual(joint, pos=pos, quat=quat, reg_target=x_prev)
        jac_target = lambda joint, r: self.ik_jac(joint, r, pos=pos, quat=quat)

        joint, _ = minimize.least_squares(
            j0,
            ik_target,
            self.bounds,
            jacobian=jac_target,
            verbose=0,
            check_derivatives=False,
            max_iter=self.max_iter,
        )
        return joint

    def solve_ik_xyz_rpy(self, xyz=None, rpy=None, j0=None, x_prev=None):
        """
        Solve the IK problem.
        pos: target position for the end effector.
        rpy: row, pitch, yaw of the end effector (in radians)
        j0: initial guess of the joint angles. If None, use the home position.
        """
        quat = Rotation.from_euler("xyz", rpy).as_quat(scalar_first=True)  # (w, x, y, z)
        return self.solve_ik_xquat(pos=xyz, quat=quat, j0=j0, x_prev=x_prev)
